package IC.asm.RegAllocation;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;

import IC.asm.ASMEntry;
import IC.asm.Instruction;
import IC.asm.Instruction.Ops;
import IC.asm.decl.CommentEntry;
import IC.asm.operands.*;
import IC.asm.operands.MachineReg.Regs;
import IC.asm.printer.AttPrinter;


public class SimpleAllocator extends DefaultAllocator<Integer,Operand> {

	private static final int REG_LVALUE = 1;
	private static final int REG_RVALUE = 2;
	
	
	public static List<ASMEntry> Process(List<ASMEntry> asmList) {
		SimpleAllocator ra = new SimpleAllocator();
		for(ASMEntry e : asmList) {
			e.accept(ra, null);
		}
		return code;
	}
	
	private Map<VirtualReg,MachineReg> 	regLocations;
	private Map<VirtualReg,Integer>		regStatus;
	private List<MachineReg> 			pool;
	private int 						poolIndex;
	private int							counter;
	
	private SimpleAllocator() {
		pool = new ArrayList<MachineReg>();
		pool.add(new MachineReg(Regs.edi));
		pool.add(new MachineReg(Regs.esi));
		pool.add(new MachineReg(Regs.ebx));
		pool.add(new MachineReg(Regs.ecx));
		counter=0;
	}

	@Override
	public Operand visit(VirtualReg that, Integer d) {
		that.bind();
		if (regLocations.get(that)!=null) {
			return regLocations.get(that);
		}
		else {
			assert(poolIndex>=0);
			MachineReg r = pool.get(poolIndex--);
			regStatus.put(that, d);
			regLocations.put(that, r);
			return r;
		}
	}
	
	@Override
	public Operand visit(MemOperand that, Integer d) {
		return new MemOperand(
			(RegOperand)that.getBase().accept(this,REG_RVALUE), 
			that.getOffset(),
			(that.getIndex() != null)?
					(RegOperand)that.getIndex().accept(this, REG_RVALUE):null,
			that.getScale());
	}
	
/**
 * 1. find all virtual registers in the instruction
 * 2. map them to real registers, on per-instruction basis
 * 3. load: mem->real regs, then inst, then save dest regs to mem.
 * 4. convert them to stack memory locations (e.g tmpi --> [-4i-4](%ebp) ) 
 * 
 */
	public Operand visit(Instruction that, Integer d) {


		regLocations = new HashMap<VirtualReg, MachineReg>();
		regStatus = new HashMap<VirtualReg, Integer>();
		// use all regs
		poolIndex = pool.size()-1;

		// find all register that are used
		Operand newSrc, newDst;
		
		newSrc = (that.getSource() != null)?that.getSource().accept(this, REG_RVALUE):null;
		
		int flags = 0;
		if (Ops.DstLvalOps.contains(that.getOp())) flags |= REG_LVALUE;
		if (Ops.DstRvalOps.contains(that.getOp())) flags |= REG_RVALUE;
		newDst = (that.getDest() != null)?that.getDest().accept(this, flags):null;

		addCode(new CommentEntry(" "+(++counter)+": "+(new AttPrinter()).visit(that, null)));
		for (VirtualReg r: regLocations.keySet()) {
			if ((regStatus.get(r) & REG_RVALUE) > 0) {
				addCode(new Instruction(Ops.MOV, r, regLocations.get(r)));
			}
		}
		addCode(new Instruction(that.getOp(), newSrc, newDst));
		for (VirtualReg r: regLocations.keySet()) {
			if ((regStatus.get(r) & REG_LVALUE) > 0) {
				addCode(new Instruction(Ops.MOV, regLocations.get(r), r));
			}
		}
		return null;
	}
}
